import threading, socket, Queue, logging
# TODO
# - Threads stay alive after disconnect?

DECIMAL, NEWLINE = range(2)

def newline_encode(msg):
  assert '\n' not in msg
  return msg + '\n'

def newline_decode(msg):
  if '\n' in msg:
    head, tail = msg.split('\n', 1)
  else:
    head, tail = None, msg
  return head, tail

def decimal_hdr_encode(msg):
  return '%d:' % len(msg) +  msg

def decimal_hdr_decode(msg):
  if ':' in msg:
    head, tail = msg.split(':', 1)
    msg_len = int(head)
    if len(tail) >= msg_len:
      return tail[:msg_len], tail[msg_len:]
  elif len(msg) > 10:
    raise DecodeError('Message header too long: %s' % msg)
  elif len(msg) > 0 and not msg.isdigit():
    raise DecodeError('Only digits are alllowed in message header: %s' % msg)
  return None, msg

class DecodeError(Exception): pass
class ConnectionError(Exception): pass

class MessageSocket:
  '''
  A MessageSocket is basically a socket, with a send and a receive queue.
  Two threads -- a reader and a writer thread -- pass messages from the socket
  in the queue, and send messages from the queue to the socket.

  If something fails, the threads should die after a while.
  '''
  def __init__(self, sock, protocol, send_qsize=10, recv_qsize=10):
    assert sock.gettimeout() == None, 'Expecting blocking socket!'
    self.sock = sock
    if protocol == DECIMAL:
      self.encode_f, self.decode_f = decimal_hdr_encode, decimal_hdr_decode
    elif protocol == NEWLINE:
      self.encode_f, self.decode_f = newline_encode, newline_decode
    else:
      raise ValueError('Unknown protocol')
    self.sendq = Queue.Queue(maxsize=send_qsize)
    self.recvq = Queue.Queue(maxsize=recv_qsize)
    self.start_threads()

  def start_threads(self):
    def fire_thread(f):
      t = threading.Thread(target=f)
      t.setDaemon(True)
      t.start()
      return t
    
    self.read_thread = fire_thread(self.__read_fun)
    self.write_thread = fire_thread(self.__write_fun)

  def test_alive(self):
    rt = self.read_thread
    wt = self.write_thread
    if not rt.is_alive():
      raise ConnectionError('Reading thread is not alive!')
    if not wt.is_alive():
      raise ConnectionError('Write thread is not alive!')

  def send(self, msg):
    self.test_alive()
    self.sendq.put_nowait(msg)

  def get(self):
    '''Return completely received messages in a list, non-blocking.'''
    self.test_alive()
    messages = []
    try:
      while True:
        messages.append(self.recvq.get_nowait())
    except Queue.Empty:
      pass
    return messages

  def get_blocking(self, timeout=None):
    '''
    Block until a message is received, or timeout seconds have passed.
    None is returned in case of a timeout
    '''
    self.test_alive()
    try:
      return self.recvq.get(timeout=timeout)
    except Queue.Empty:
      return None

  def __write_fun(self):
    try:
      while True:
        try:
          msg = self.sendq.get(True, 10)
          self.sock.sendall(self.encode_f(msg))
        except Queue.Empty:
          pass
        self.test_alive()
          
    except Exception as e:
      self.log.warning(e)
      self.sock.close()
      
  def __read_fun(self):
    buff = ''
    try:
      while True:
        s = self.sock.recv(1024)
        if s == '':
          raise IOError('The socket has been closed, recv return 0 bytes.')
        buff += s
        while True:
          msg, buff = self.decode_f(buff)
          if msg == None: 
            break
          self.recvq.put(msg)
    except Exception as e:
      self.log.warning(e)
      self.sock.close()

  @property
  def log(self):
    return logging.getLogger('tinynl.MessageSocket')

class Server:
  def __init__(self, port=None, protocol=DECIMAL):
    '''
    Start server thread that waits for connections. When a connection is made,
    a MessageSocket is created to wrap the socket, and is added to the client
    queue.
    '''
    self.protocol = protocol
    # Create server socket
    ss = self.serv_sock = socket.socket()
    ss.bind(('', port if port else 0))
    self.port = ss.getsockname()[1]
    ss.listen(5)
    self.log.info('Listening on port %d.' % self.port)

    # Create client Queue
    self.clientq = Queue.Queue(maxsize=50)

    # Start worker thread.
    t = threading.Thread(target=self.worker)
    t.setDaemon(True)
    t.start()
    self.work_thread = t

  def worker(self):
    '''
    Listens on server socket, accepts new connections, and add a
    MessageSocket to the client list.
    '''
    while True:
      sock, addr = self.serv_sock.accept()
      self.clientq.put(MessageSocket(sock, self.protocol))
      self.log.info('Added client (%s:%d)' % addr)

  def new_clients(self):
    '''
    Convenience function to get new clients from the client queue.
    Returns a list, with zero or more MessageSockets.
    Handling the clients is the responsibility of the user.
    '''
    result = []
    while not self.clientq.empty():
      result.append(self.clientq.get())
    return result
  
  @property
  def log(self):
    return logging.getLogger('tinynl.Server')

class Client(MessageSocket):
  def __init__(self, host, port, protocol=DECIMAL):
    sock = socket.socket()
    sock.connect((host, port))
    MessageSocket.__init__(self, sock, protocol)

